import time
from functools import wraps


def func_timer(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        start = time.time()
        result = func(*args, **kwargs)
        end = time.time()
        print(f"{func.__name__} runtime: {end - start:.6f} second")
        return result

    return wrapper


global time_dict, time_id

time_dict = {}
time_id = 0


def tick():
    global time_id
    time_dict[time_id] = time.time()
    time_id += 1
    return time_id - 1


def tock(prev_id):
    ans = time.time() - time_dict[prev_id]
    del time_dict[prev_id]
    return ans


class Timer:
    def __init__(self, seg_description: str):
        self.desc = seg_description

    def __enter__(self):
        self.start = time.time()

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.end = time.time()
        print(f"[{self.desc}] took {self.end - self.start:.6f} seconds")


@func_timer
def test():
    time.sleep(2)


if __name__ == '__main__':
    with Timer('test'):
        test()
